package com.fanshuai.websocket.server;

import com.fanshuai.websocket.annotation.WebSocketMapping;
import org.springframework.beans.BeansException;
import org.springframework.context.ApplicationContext;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.AnnotationUtils;
import org.springframework.web.reactive.handler.SimpleUrlHandlerMapping;
import org.springframework.web.reactive.socket.WebSocketHandler;

import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

public class WebSocketHandlerMapping extends SimpleUrlHandlerMapping {
    private Map<String, WebSocketHandler> handlerMap = new HashMap<>();

    @Override
    public void initApplicationContext() throws BeansException {
        ApplicationContext context = obtainApplicationContext();

        Map<String, Object> beanMap = context.getBeansWithAnnotation(WebSocketMapping.class);
        beanMap.values().forEach(bean -> {
            if (!(bean instanceof WebSocketHandler)) {
                throw new RuntimeException("@WebSocketMapping is used in non WebSocketHandler class");
            }

            WebSocketMapping webSocketMapping = AnnotationUtils.getAnnotation(bean.getClass(), WebSocketMapping.class);
            handlerMap.put(Objects.requireNonNull(webSocketMapping).value(), (WebSocketHandler) bean);
        });
        super.setOrder(Ordered.HIGHEST_PRECEDENCE);
        super.setUrlMap(handlerMap);

        super.initApplicationContext();
    }
}
